package 中等.搜索.回溯;

import java.util.Arrays;

/**
 * 给定一个整数数组  nums 和一个正整数 k，找出是否有可能把这个数组分成 k 个非空子集，其总和都相等。
 * <p>
 * 来源：力扣（LeetCode）
 * 链接：https://leetcode.cn/problems/partition-to-k-equal-sum-subsets/
 */
public class 划分为k个相等的子集_698 {

    public static void main(String[] args) {

        划分为k个相等的子集_698 obj = new 划分为k个相等的子集_698();
        System.out.println(obj.canPartitionKSubsets(new int[]{1, 2, 1, 2, 1, 2}, 3));

    }

    /**
     * 回溯+剪枝
     * 和 火柴拼正方形_473 的思路一致
     * 创建 k 个容量为 sum/k 的桶，尝试将 nums 中所有元素放入桶中
     *
     * @param nums
     * @param k
     * @return
     */
    public boolean canPartitionKSubsets(int[] nums, int k) {
        int sum = 0;
        for (int num : nums) {
            sum += num;
        }
        if (sum % k != 0) return false;

        Arrays.sort(nums);
        // 最小的都比平均值大，不可能符合条件
        if (nums[0] > sum / k) return false;

        int[] buckets = new int[k];
        Arrays.fill(buckets, sum / k);

        // 优先放置大的值，使得快速失败
        return dfs(nums.length - 1, buckets, nums);
    }

    // 尝试将 nums[idx] 放入到 buckets 某个桶中
    private boolean dfs(int idx, int[] buckets, int[] nums) {
        // 放完了
        if (idx == -1) return true;

        for (int i = 0; i < buckets.length; i++) {
            // 剪枝
            if (buckets[i] < nums[idx] || i > 0 && buckets[i] == buckets[i - 1]) {
                continue;
            }

            buckets[i] -= nums[idx];
            if (dfs(idx - 1, buckets, nums)) {
                return true;
            }
            buckets[i] += nums[idx];
        }
        return false;
    }

}
